#!/usr/bin/env python3
"""
系统设置脚本
用于初始化数据库、下载模型等
"""

import asyncio
import os
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.core.config import config


async def setup_neo4j():
    """设置Neo4j数据库"""
    print("🔧 Setting up Neo4j database...")
    
    try:
        from src.data_layer.graph_store.neo4j_store import Neo4jGraphStore
        
        graph_config = config.get("data_layer.graph_store.config", {})
        graph_store = Neo4jGraphStore(**graph_config)
        
        # 创建索引和约束
        await graph_store.setup_schema()
        print("✅ Neo4j setup completed!")
        
    except Exception as e:
        print(f"❌ Neo4j setup failed: {e}")


async def setup_vector_db():
    """设置向量数据库"""
    print("🔧 Setting up vector database...")
    
    try:
        from src.data_layer.vector_store.chroma_store import ChromaVectorStore
        
        vector_config = config.get("data_layer.vector_store.config", {})
        vector_store = ChromaVectorStore(**vector_config)
        
        # 初始化集合
        await vector_store.initialize()
        print("✅ Vector database setup completed!")
        
    except Exception as e:
        print(f"❌ Vector database setup failed: {e}")


async def download_models():
    """下载必要的模型"""
    print("📥 Downloading required models...")
    
    # 下载嵌入模型
    try:
        from sentence_transformers import SentenceTransformer
        
        embedding_model = config.get("lightrag_engine.embedding.model")
        if embedding_model:
            print(f"Downloading embedding model: {embedding_model}")
            SentenceTransformer(embedding_model)
            print("✅ Embedding model downloaded!")
            
    except Exception as e:
        print(f"❌ Failed to download embedding model: {e}")
    
    # 下载重排模型
    try:
        rerank_model = config.get("lightrag_engine.reranking.model")
        if rerank_model:
            print(f"Downloading reranking model: {rerank_model}")
            SentenceTransformer(rerank_model)
            print("✅ Reranking model downloaded!")
            
    except Exception as e:
        print(f"❌ Failed to download reranking model: {e}")


def create_directories():
    """创建必要的目录"""
    print("📁 Creating directories...")
    
    directories = [
        "data/raw",
        "data/processed", 
        "data/vector_db",
        "data/graph_db",
        "logs",
        "models"
    ]
    
    for directory in directories:
        Path(directory).mkdir(parents=True, exist_ok=True)
        
    print("✅ Directories created!")


def check_environment():
    """检查环境变量"""
    print("🔍 Checking environment variables...")
    
    required_vars = [
        "OPENAI_API_KEY",
        "NEO4J_URI", 
        "NEO4J_USERNAME",
        "NEO4J_PASSWORD"
    ]
    
    missing_vars = []
    for var in required_vars:
        if not os.getenv(var):
            missing_vars.append(var)
    
    if missing_vars:
        print("⚠️  Missing environment variables:")
        for var in missing_vars:
            print(f"  - {var}")
        print("Please set these variables in your .env file")
        return False
    
    print("✅ Environment variables check passed!")
    return True


async def main():
    """主设置流程"""
    print("🚀 Starting BoLe HR Platform setup...\n")
    
    # 检查环境
    if not check_environment():
        sys.exit(1)
    
    # 创建目录
    create_directories()
    
    # 设置数据库
    await setup_vector_db()
    await setup_neo4j()
    
    # 下载模型
    await download_models()
    
    print("\n🎉 Setup completed successfully!")
    print("You can now run: python main.py serve")


if __name__ == "__main__":
    asyncio.run(main())